import numpy as np
import matplotlib.pyplot as plt
import pywt
from magnet_wave_signal import build_signal

class WaveletFeatureAnalyzer:
    """
    小波特征分析器，用于捕捉信号的平滑波形特征
    """
    
    def __init__(self, wavelets=['sym2', 'sym5', 'coif1', 'coif3']):
        """
        初始化分析器
        参数：
            wavelets: 用于分析的小波列表
        """
        self.wavelets = wavelets
        
    def _analyze_coefficients(self, coeffs, wavelet_name):
        """
        分析小波系数的符号序列
        参数：
            coeffs: 小波分解系数
            wavelet_name: 小波名称
        返回：
            分析结果字典
        """
        cA, cD = coeffs
        
        # 分析细节系数的符号变化
        sign_changes = np.diff(np.sign(cD))
        positive_peaks = np.where((cD[:-1] > 0) & (cD[1:] < 0))[0]
        negative_peaks = np.where((cD[:-1] < 0) & (cD[1:] > 0))[0]
        
        # 计算最大正负值及其位置
        max_pos_idx = np.argmax(cD)
        max_neg_idx = np.argmin(cD)
        
        return {
            'wavelet': wavelet_name,
            'cA_mean': np.mean(cA),
            'cD_max': np.max(cD),
            'cD_min': np.min(cD),
            'max_pos_idx': max_pos_idx,
            'max_neg_idx': max_neg_idx,
            'sign_changes': len(positive_peaks) + len(negative_peaks),
            'pos_neg_pattern': max_pos_idx < max_neg_idx  # True表示先正后负
        }
    
    def analyze_signal(self, signal):
        """
        对信号进行多小波分析
        参数：
            signal: 输入信号
        返回：
            分析结果列表
        """
        results = []
        
        for wavelet in self.wavelets:
            # 进行单层小波分解
            coeffs = pywt.dwt(signal, wavelet)
            analysis = self._analyze_coefficients(coeffs, wavelet)
            results.append(analysis)
            
        return results
    
    def _plot_wavelet_analysis(self, signal, results):
        """
        绘制小波分析结果
        参数：
            signal: 原始信号
            results: 分析结果
        """
        fig, axes = plt.subplots(len(self.wavelets) + 1, 2, figsize=(12, 2*(len(self.wavelets)+1)))
        
        # 绘制原始信号
        axes[0, 0].plot(signal, 'k', lw=2)
        axes[0, 0].set_title('原始信号')
        axes[0, 0].grid(True, alpha=0.3)
        axes[0, 1].axis('off')
        
        # 绘制各小波的分解结果
        for i, (wavelet, result) in enumerate(zip(self.wavelets, results)):
            coeffs = pywt.dwt(signal, wavelet)
            cA, cD = coeffs
            
            # 绘制近似系数
            axes[i+1, 0].plot(cA, 'b-', lw=1.5, label='近似系数 (cA)')
            axes[i+1, 0].set_title(f'{wavelet} - 近似系数')
            axes[i+1, 0].grid(True, alpha=0.3)
            axes[i+1, 0].legend()
            
            # 绘制细节系数并标记特征
            axes[i+1, 1].plot(cD, 'r-', lw=1.5, label='细节系数 (cD)')
            
            # 标记最大正值和最大负值
            max_pos_idx = result['max_pos_idx']
            max_neg_idx = result['max_neg_idx']
            
            axes[i+1, 1].plot(max_pos_idx, cD[max_pos_idx], 'go', markersize=8, label='最大正值')
            axes[i+1, 1].plot(max_neg_idx, cD[max_neg_idx], 'ro', markersize=8, label='最大负值')
            
            # 添加零线
            axes[i+1, 1].axhline(y=0, color='k', linestyle='--', alpha=0.5)
            
            pattern = "先正后负" if result['pos_neg_pattern'] else "先负后正"
            axes[i+1, 1].set_title(f'{wavelet} - 细节系数 ({pattern})')
            axes[i+1, 1].grid(True, alpha=0.3)
            axes[i+1, 1].legend()
        plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 更好的中文字体，可以显示“-”号
        plt.tight_layout()
        plt.show()
    
    def detect_pattern(self, signal, threshold=0.1):
        """
        检测信号中的"先上后下"模式
        参数：
            signal: 输入信号
            threshold: 检测阈值
        返回：
            检测结果
        """
        results = self.analyze_signal(signal)
        
        # 统计各小波的模式检测结果
        pattern_votes = sum(1 for r in results if r['pos_neg_pattern'])
        confidence = pattern_votes / len(results)
        
        print("=== 小波特征分析结果 ===")
        for result in results:
            pattern = "先正后负" if result['pos_neg_pattern'] else "先负后正"
            print(f"{result['wavelet']:>6}: 最大正值位置={result['max_pos_idx']:>3}, "
                  f"最大负值位置={result['max_neg_idx']:>3}, 模式={pattern}")
        
        print(f"\n模式检测置信度: {confidence:.2%} ({pattern_votes}/{len(results)} 个小波支持'先正后负'模式)")
        
        # 绘制分析结果
        self._plot_wavelet_analysis(signal, results)
        
        return {
            'pattern_detected': confidence > 0.5,
            'confidence': confidence,
            'details': results
        }

def _create_test_signals():
    """
    创建测试信号：包括"先上后下"和"先下后上"两种模式
    """
    # 原始的"先上后下"信号
    x1, signal1 = build_signal()
    
    # 创建"先下后上"信号（翻转原信号）
    signal2 = -signal1
    
    return signal1, signal2

if __name__ == '__main__':
    # 创建分析器
    analyzer = WaveletFeatureAnalyzer()
    
    # 生成测试信号
    signal_up_down, signal_down_up = _create_test_signals()
    
    print("分析 '先上后下' 信号:")
    result1 = analyzer.detect_pattern(signal_up_down)
    
    print("\n" + "="*50)
    print("分析 '先下后上' 信号:")
    result2 = analyzer.detect_pattern(signal_down_up)